#!/usr/bin/env python3
# -*- encoding utf-8 -*-

'''
@File: save_mnist_to_jpg.py
@Date: 2024-08-23
@Author: KRISNAT
@Version: 0.0.0
@Email: ****
@Copyright: (C)Copyright 2024, KRISNAT
@Desc:
    1. 通过 torchvision.datasets.MNIST 下载、解压和读取 MNIST 数据集；
    2. 使用 PIL.Image.save 将 MNIST 数据集中的灰度图片以 JPEG 格式保存。
'''

import sys, os
sys.path.insert(0, os.getcwd())

from torchvision.datasets import MNIST
import PIL
from tqdm import tqdm

def get_data(root):
    # 从网络上下载或从本地加载MNIST数据集
    # 训练集60K、测试集10K
    # torchvision.datasets.MNIST接口下载的数据一组元组
    # 每个元组的结构是: (PIL.Image.Image image model=L size=28x28, 标签数字 int)
    training_dataset = MNIST(
        root=root,
        train=True,
        download=False,
    )
    test_dataset = MNIST(
        root=root,
        train=False,
        download=False,
    )
    return training_dataset, test_dataset

def save2jpeg(training_dataset, test_dataset):
    # 保存训练集图片
    with tqdm(total=len(training_dataset), ncols=150) as pro_bar:
        for idx, (X, y) in enumerate(training_dataset):
            f = root + "/" + "training_" + str(idx) + \
                "_" + str(training_dataset[idx][1] )+ ".jpg"  # 文件路径
            training_dataset[idx][0].save(f)
            pro_bar.update(n=1)

    # 保存测试集图片
    with tqdm(total=len(test_dataset), ncols=150) as pro_bar:
        for idx, (X, y) in enumerate(test_dataset):
            f = root + "/" + "test_" + str(idx) + \
                "_" + str(test_dataset[idx][1] )+ ".jpg"  # 文件路径
            test_dataset[idx][0].save(f)
            pro_bar.update(n=1)


def show(dataset):
    # 可视化
    import matplotlib.pyplot as plt
    plt.rcParams['font.sans-serif'] = ['SimHei']
    fig, ax = plt.subplots(ncols=3, nrows=3)
    ax[0, 0].imshow(dataset[0][0], cmap=plt.cm.gray)
    ax[0, 0].set_title(f"标签为{dataset[0][1]}")
    ax[0, 1].imshow(dataset[1][0], cmap=plt.cm.gray)
    ax[0, 1].set_title(f"标签为{dataset[1][1]}")
    ax[0, 2].imshow(dataset[2][0], cmap=plt.cm.gray)
    ax[0, 2].set_title(f"标签为{dataset[2][1]}")
    ax[1, 0].imshow(dataset[3][0], cmap=plt.cm.gray)
    ax[1, 0].set_title(f"标签为{dataset[3][1]}")
    ax[1, 1].imshow(dataset[4][0], cmap=plt.cm.gray)
    ax[1, 1].set_title(f"标签为{dataset[4][1]}")
    ax[1, 2].imshow(dataset[5][0], cmap=plt.cm.gray)
    ax[1, 2].set_title(f"标签为{dataset[5][1]}")
    ax[2, 0].imshow(dataset[5][0], cmap=plt.cm.gray)
    ax[2, 0].set_title(f"标签为{dataset[5][1]}")
    ax[2, 1].imshow(dataset[6][0], cmap=plt.cm.gray)
    ax[2, 1].set_title(f"标签为{dataset[6][1]}")
    ax[2, 2].imshow(dataset[7][0], cmap=plt.cm.gray)
    ax[2, 2].set_title(f"标签为{dataset[7][1]}")
    plt.xticks([])
    plt.yticks([])
    plt.show()  # 显示绘图

if __name__ == "__main__":
    # 图片保存路径
    root = 'mnist_jpg'
    if not os.path.exists(root):
        os.makedirs(root)

    training_dataset, test_dataset = get_data("./dataset")



    # save2jpeg(training_dataset, test_dataset)

    show(training_dataset)

